# import matplotlib
# matplotlib.use('Agg')
# import copy
 
# from utils.options import args_parser
# from models.Update import *
# from models.resnet import *
# from models.test import *
# from models.aggregation import *
# from models.branchnet import *
# from algorithm.Training_FedAvg import *
# from algorithm.Training_FedClass import *
# from algorithm.Training_FedClass2 import *
# from utils.get_dataset import get_dataset
# from utils.save_result import save_result
# from utils.set_seed import set_random_seed



# if __name__ == '__main__':
#     args = args_parser()
#     args.device = torch.device('cuda:{}'.format(args.device) if torch.cuda.is_available() and args.device != -1 else 'cpu')
#     print(torch.cuda.is_available())
#     set_random_seed(args.seed)

#     dataset_train,dataset_test,dict_users ,dict_global= get_dataset(args)
#     net_list = []
#     # net_num = args.num_groups
#     # for i in range(net_num):
#     #     net_list.append(ResNet18_cifar10(num_classes = args.num_classes))
#     # print(ResNet50_cifar10(num_classes = args.num_classes))
#     # net_list.append(ResNet8(num_classes = args.num_classes))
#     # net_list.append(ResNet18_cifar10(num_classes = args.num_classes))
#     # net_list.append(ResNet18_cifar10(num_classes = args.num_classes))
#     # net_list.append(ResNet34_cifar10(num_classes = args.num_classes))
    
#     # for id in range(args.num_groups):
#     #     net_list[id].to(args.device)

#     # print(net_list)
#     print(args.algorithm + '_' + args.dataset+'_'+args.model+'_'+str(args.iid)+'_'+str(args.noniid_case)+'_'+str(args.data_beta))
#     if args.algorithm == 'FedTest':
#         FedTest(args,net_list,dataset_train,dataset_test,dict_users,dict_global)
#     elif args.algorithm == 'FedAvg':
#         net_glob = ResNet18_cifar10(num_classes = args.num_classes)
#         net_glob.to(args.device)
#         FedAvg(args,net_glob,dataset_train,dataset_test,dict_users)
#     elif args.algorithm == 'FedHomo':
#         net_list[1].to(args.device)
#         FedHomo(args,net_list[1],dataset_train,dataset_test,dict_users,dict_global)
#     elif args.algorithm == 'Fedpre':
#         Fedpre(args,net_list,dataset_train,dataset_test,dict_users,dict_global)
#     elif args.algorithm == 'Fednon':
#         Fednon(args,net_list,dataset_train,dataset_test,dict_users,dict_global)
#     elif args.algorithm == 'Fedloss':
#         # for name,p in net_list[0].named_parameters():
#             # if "layer1" in name:
#             # print(p)
#         # print(net_list[0].named_parameters())
#         # models = []
#         # for i in range (args.num_groups):
#         #     models.append([])
#         # for i in range (args.num_groups):
            
#         #     model = net_list[i].to(args.device)
#         #     # args.models[i].append(copy.deepcopy(model))
#         #     for j in range (args.num_block):
#         #         models[i].append(copy.deepcopy(model))
#         #         for name,p in models[i][j].named_parameters():
#         #             p.requires_grad = False
#         #             for layer in layers0:
#         #                 if layer == name:
#         #                     p.requires_grad = True
#         #         for name,p in models[i][j].named_parameters(): 
#         #             for layer in layers[j]:
#         #                 if layer in name:
#         #                     p.requires_grad = True
#         # for j in range(args.num_block):
#         #     for name,p in models[0][j].named_parameters():
#         #         if p.requires_grad == True:
#         #             print(name)
#         #     print('*'*15)

#         Fedloss(args,net_list,dataset_train,dataset_test,dict_users,dict_global)
    
#     elif args.algorithm == "Fedtwo":
#         net_list_1 = []
#         for i in range(args.num_groups):
#             net_list_1.append(ResNet18_cifar10(num_classes = args.num_classes))
#         net_list_2 = []
#         for i in range(args.num_groups):
#             net_list_2.append(ResNet34_cifar10(num_classes = args.num_classes))
#         for id in range(args.num_groups):
#             net_list_1[id].to(args.device)
#             net_list_2[id].to(args.device)
#         Fedtwo(args,net_list_1,net_list_2,dataset_train,dataset_test,dict_users,dict_global)
    
#     elif args.algorithm == "Fedglob":
#         net_list_1 = []
#         for i in range(args.num_groups):
#             net_list_1.append(ResNet18_cifar10(num_classes = args.num_classes))
#         net_list_2 = []
#         for i in range(args.num_groups):
#             net_list_2.append(ResNet34_cifar10(num_classes = args.num_classes))
#         for id in range(args.num_groups):
#             net_list_1[id].to(args.device)
#             net_list_2[id].to(args.device)
#         Fedglob(args,net_list_1,net_list_2,dataset_train,dataset_test,dict_users,dict_global)

#     elif args.algorithm == "FedC":
#         net_glob = ResNet18_cifar10(num_classes = args.num_classes)
#         net_glob.to(args.device)
#         FedC(args,net_glob,dataset_train,dataset_test,dict_users,dict_global)

#     elif args.algorithm == "FedClass":
#         net_glob = ResNet18_cifar10(num_classes = args.num_classes)
#         net_glob.to(args.device)
#         FedClass(args,net_glob,dataset_train,dataset_test,dict_users,dict_global)
#     elif args.algorithm == "FedClass2":
#         net_glob = ResNet18_cifar10(num_classes = args.num_classes)
#         net_glob.to(args.device)
#         FedClass2(args,net_glob,dataset_train,dataset_test,dict_users,dict_global)

        



